# -*- coding: utf-8 -*-
"""
@Time    : 2024/7/22 10:06 
@Author  : ZhangShenao 
@File    : in_memory_vector_store.py
@Desc    : 内存向量数据库

实现自定义的向量数据库,需要继承VectorStore基类,并重写similarity_search和from_texts方法
"""
import uuid
from typing import Type, List, Optional, Any, Iterable

import numpy as np
from langchain_core.documents import Document
from langchain_core.embeddings import Embeddings
from langchain_core.vectorstores import VectorStore, VST


class InMemoryVectorStore(VectorStore):
    """
    内存向量数据库
    将全部向量数据存储在内存中
    采用欧氏距离计算相似度
    """

    def __init__(self, embedding: Embeddings):
        """
        构造函数
        :param embedding: Embedding模型
        """

        # 初始化Embedding模型
        self.__embedding = embedding

        # 初始化内存向量存储
        # key为文档ID, value为文档记录(包括文本、向量和元数据等)
        self.__doc_store = {}

    def similarity_search(self, query: str, k: int = 4, **kwargs: Any) -> List[Document]:
        # 将queryEmbedding成向量
        query_vector = self.__embedding.embed_query(query)

        # 变量所有向量,计算欧式记录
        result = []
        for doc_id, doc in self.__doc_store.items():
            doc_vector = doc['vector']
            distance = self._euclidean_distance(query_vector, doc_vector)
            result.append(
                {'distance': distance,
                 **doc,  # 将doc中的key-value结构,合并至result中
                 }
            )

        # 对result中的distance升序排序
        sorted_result = sorted(result, key=lambda x: x['distance'], reverse=False)

        # 获取距离最近的k条记录
        k_nearest = sorted_result[:k]

        # 构造Document记录并返回
        documents = [Document(page_content=doc['text'], metadata={**doc['metadata'], "score": doc['distance']})
                     for doc in k_nearest
                     ]
        return documents

    def add_texts(
            self,
            texts: Iterable[str],
            metadatas: Optional[List[dict]] = None,
            # One of the kwargs should be `ids` which is a list of ids
            # associated with the texts.
            # This is not yet enforced in the type signature for backwards compatibility
            # with existing implementations.
            **kwargs: Any,
    ) -> List[str]:

        # 校验texts和metadatas
        if len(list(texts)) == 0:
            raise ValueError('texts must no be empty!')
        if metadatas is not None and len(list(texts)) != len(metadatas):
            raise ValueError('texts and metadatas must have the same length!')

        # 将文本Embedding成向量
        embeddings = self.__embedding.embed_documents(list(texts))

        # 使用UUID,生成文档ID
        doc_ids = [str(uuid.uuid4()) for text in texts]

        # 构造文档列表,并保存至内存中
        for i, text in enumerate(texts):
            doc_id = doc_ids[i]
            vector = embeddings[i]
            doc = {
                'id': doc_id,
                'text': text,
                'vector': vector,
                'metadata': metadatas[i] if metadatas is not None else {}
            }
            self.__doc_store[doc_id] = doc

        # 返回生成的文档ID
        return doc_ids

    @classmethod
    def from_texts(cls: Type[VST], texts: List[str], embedding: Embeddings, metadatas: Optional[List[dict]] = None,
                   **kwargs: Any) -> VST:

        # 构造向量数据库实例
        vector_store = cls(embedding=embedding)

        # 添加文档
        vector_store.add_texts(texts=texts, metadatas=metadatas, **kwargs)

        # 返回向量数据库实例
        return vector_store

    @classmethod
    def _euclidean_distance(cls, vector1: List[float], vector2: List[float]) -> float:
        """
        计算两个向量的欧氏距离
        :param vector1: 向量1
        :param vector2: 向量2
        :return: 欧氏距离
        """

        # 使用Numpy包直接计算欧氏距离
        return np.linalg.norm(np.array(vector1) - np.array(vector2))
